#!/usr/bin/python



import popen2
import os,sys,re
from time import time,sleep
import signal
import glob

from pdfobj import *

class pdf:
    DEBUG = 0
    def __init__(self,indata,infile):
        self.indata = indata
        self.infile = infile
        self.objects = {} 
        self.pages = []

    def parse(self):     
        trailers = re.findall('trailer\s*(.*?)\s*startxref',self.indata,re.MULTILINE|re.DOTALL)
        if not trailers: trailers = re.findall('trailer\s*(<<.*?>>)\s*',self.indata,re.DOTALL)   #some samples don't have the key words "startxref"
        if trailers:
           print "reg find trailer"
           for trailer in trailers:
               self.objects["trailer"] = pdfobj("trailer",trailer)  #asign the last trailer to self.objects["trailer"]
               print "self.objects[\"trailer\"].tags:\n%s"%self.objects["trailer"].tags
     

        objs = re.findall('\n?(\d+)\s+(\d+)\s+obj[\s]*(.*?)\s*\n?(endobj|objend)',self.indata,re.MULTILINE|re.DOTALL)
        if objs:
            for obj in objs:
                #fill all objects
                key = obj[0] + ' ' + obj[1]
                self.objects[key] = pdfobj(key,obj[2])
 

    def is_valid(self):
        #if self.indata.startswith('%PDF-') or self.indata.startswith('%%PDF-'):
            #return True
        #return False
        return True
        
    def __repr__(self):
        if not self.is_valid():
            return 'Invalid PDF file "%s"' % (self.infile)
        out = 'PDF file %s has %d obj items\n' % (self.infile,len(self.objects))
        for obj in sorted(self.objects.keys()):
            out += str(self.objects[obj]) + '\n'

        return out
    
    def check_jbig2(self):
        print "in check_jbig2"
        for indirect_refer in self.objects.keys():
            print "for indirect_refer in self.objects.keys():, indirect_refer:%s"%indirect_refer
            indirect_obj = self.objects[indirect_refer]
            if indirect_obj.obj_check_jbig2() == True:
               print "//[alert] cve-2009-0658, in this file,obj %s have jbig2 attack"
               fout = open(self.infile+".jbig2_alert",'a')
               fout.write("//[alert] cve-2009-0658, in this file,obj %s have jbig2 attack"%indirect_refer)
               fout.close()
               

    def getJavaScript(self):
        print "in line 630, in getJavaScript function"
        have_trailer = False
        have_catalog = False
        
        if "trailer" in self.objects.keys():
            trailerobj = self.objects["trailer"]
            have_trailer = True
            print "have_trailer = True"
            if "Root" in self.objects["trailer"].tags.keys():
               have_catalog = True
               catalognum = self.objects["trailer"].tags["Root"][1]
               catalogobj = self.objects[catalognum]
               print "have_catalog = True"
        out= ""  #tmp argument
        js_content = ""
        header =""
       

        #################################################################
        ######get header through the /Info in trailer
        #################################################################
        print "in line 659, begin to get header through the /Info in trailer: \n"
        if have_trailer and "Info" in trailerobj.tags.keys():
            if trailerobj.tags["Info"][0] == "refer":
                header += self.get_header_from_information_Dictionary(self.objects[trailerobj.tags["Info"][1]])
            elif trailerobj.tags["Info"][0] == "dictionary obj":
                header += self.get_header_from_information_Dictionary(trailerobj.tags["Info"][1]) 
        
        print "in line 621, get header through the /Info in trailer: \n%s"%header


        #################################################################
        ######get header through the /Annots in pages
        #################################################################
        print "in line 681, begin to get header through the /Annots in pages: \n"
        print "catalogobj",catalogobj
        print "catalogobj.tags",catalogobj.tags
        if have_catalog and "Pages" in catalogobj.tags.keys():
            if catalogobj.tags["Pages"][0] == "refer":
                pagesObj =  self.objects[catalogobj.tags["Pages"][1]]
            elif catalogobj.tags["Pages"][0] == "dictionary obj":
                pagesObj = catalogobj.tags["Pages"][1]
            if pagesObj and "Kids" in pagesObj.tags.keys():
                if pagesObj.tags["Kids"][0] == "array": 
                    page_list =  re.findall("(\d+\s+\d+)\s+R",pagesObj.tags["Kids"][1])
                elif pagesObj.tags["Kids"][0] == "refer":
                    page_list =  [pagesObj.tags["Kids"][1]]
                header += "page_list = new Array();\n"
                for page_obj_num in page_list:
                    header += "page = new Page();\nannot_array = new Array();\n"
                    if page_obj_num not in self.objects.keys():continue
                    pageObj = self.objects[page_obj_num]
                    if pageObj and "Annots" in pageObj.tags.keys() and pageObj.tags["Annots"][0] == "array": 
                       annot_num_list =  re.findall("(\d+\s+\d+)\s+R",pageObj.tags["Annots"][1])
                       for annot_num in annot_num_list:
                           header += "annotation = new Annotation();\n"
                           header += self.get_header_from_annotation(self.objects[annot_num])
                           header += "annot_array.push(annotation);\n"
                    header += "page['annots'] = annot_array;\npage_list.push(page);\n"  

        print "in line 681, get header through the /Annots in pages: \n%s"

        #################################################################
        ######get document_level_javascript through the /Names in catalog
        #################################################################
        print "in line 670, begin to get document_level_javascript through the /Names in catalog:\n"
        out = ""
        if have_catalog and "Names" in catalogobj.tags.keys():
            if catalogobj.tags["Names"][0] == "refer":
               namedictionary = self.objects[catalogobj.tags["Names"][1]]
            elif catalogobj.tags["Names"][0] == "dictionary obj":
               namedictionary = catalogobj.tags["Names"][1]
            else:
               print "error"
            print "namedictionary: ",namedictionary            
            if namedictionary and "JavaScript" in namedictionary.tags.keys():
               print "\"JavaScript\" in namedictionary.tags.keys():"
               if namedictionary.tags["JavaScript"][0] == "refer":
                  name_tree_dictionary = self.objects[namedictionary.tags["JavaScript"][1]]
               elif namedictionary.tags["JavaScript"][0]  == "dictionary obj":
                  name_tree_dictionary = namedictionary.tags["JavaScript"][1]
               else:
                  print "error"
               if name_tree_dictionary and "Names" in name_tree_dictionary.tags.keys() and name_tree_dictionary.tags["Names"][0]=="array":
                    print "\"Names\" in name_tree_dictionary.tags.keys() and name_tree_dictionary.tags[\"Names\"][0]==\"array\":"
                    global_names  = re.findall('\((.+?)\)\s*(\d+\s+\d+)\s+R',name_tree_dictionary.tags["Names"][1],re.MULTILINE|re.DOTALL)
                    if global_names:
                       for global_names_key,global_names_value in global_names:
                          print "for global_names_key,global_names_value in global_names: global_names_value is %s"%global_names_value     
                          tmp = ""  
                          tmp = self.get_javascript_from_actions(global_names_value)          
                          out  +=  tmp
                          print "self.get_javascript_from_actions(global_names_value):%s\n"%tmp
                    else:
                      print "get global_names error"
        js_content +=  out
        print "in line 632, get document_level_javascript through the /Names in catalog:\n%s"%out     
  
        #################################################################
        ######get javascript through the /OpenAction in catalog
        #################################################################
        print "\nin line 700, begin to get javascript through the /OpenAction in catalog:\n"
        out = ""
        if have_catalog and "OpenAction" in catalogobj.tags.keys():
            if catalogobj.tags["OpenAction"][0] == "refer":
                out = self.get_javascript_from_actions(self.objects[catalogobj.tags["OpenAction"][1]])
            elif catalogobj.tags["OpenAction"][0] == "dictionary obj":
                 out = self.get_javascript_from_actions(catalogobj.tags["OpenAction"][1]) 
            #elif catalogobj["OpenAction"][0] == "array":     
        
        js_content +=  out
        print "\nin line 659, get javascript through the /OpenAction in catalog:\n%s"% out

        #################################################################
        ###### page, annot and other objects may have /AA
        #################################################################
        print "in line 719,  begin to get /AA in page, annot:\n"
        out = ""
        for indirect_refer in self.objects.keys():
            print "for indirect_refer in self.objects.keys():, indirect_refer:%s"%indirect_refer
            indirect_obj = self.objects[indirect_refer]
            if "AA" in indirect_obj.tags.keys():
                if indirect_obj.tags["AA"][0] == "refer":
                    out += self.get_javascript_from_aditional_actions(self.objects[indirect_obj.tags["AA"][1]])
                elif indirect_obj.tags["AA"][0] == "dictionary obj":
                    out += self.get_javascript_from_aditional_actions(indirect_obj.tags["AA"][1]) 
        js_content +=  out
        print "in line 677,  page, annot and other objects may have /AA :\n%s"% out
                
        return js_content,header

    def get_header_from_information_Dictionary(self, information_dictionary):
        print  "in function get_header_from_information_Dictionary"
        if isinstance(information_dictionary,str):     #information_dictionary argument is the objnum
            information_dictionary_obj = self.objects[information_dictionary]
        else:
            information_dictionary_obj = information_dictionary      #information_dictionary argument is a dictionary obj
        document_information = ""
        print information_dictionary_obj.tags
        for key in information_dictionary_obj.tags.keys():
            print "key %s in information_dictionary_obj,type: %s"%(key,information_dictionary_obj.tags[key][0])
            if information_dictionary_obj.tags[key][0] =="refer":
                   document_information += 'info.%s = String(\"%s\");\napp.doc.%s = String(\"%s\");\n' % (key.lower(),self.objects[information_dictionary_obj.tags[key][1]].decode_stream(),key.lower(),self.objects[information_dictionary_obj.tags[key][1]].decode_stream())
            elif information_dictionary_obj.tags[key][0] =="literal string":
                   document_information += 'info.%s = String(\"%s\");\napp.doc.%s = String(\"%s\");\n' % (key.lower(),get_bracketed_content(information_dictionary_obj.tags[key][1]),key.lower(),get_bracketed_content(information_dictionary_obj.tags[key][1]))
        return document_information 

    def get_javascript_from_actions(self, action):
        if isinstance(action,str):     #action argument is the objnum
           actionobj = self.objects[action]
        else:
           actionobj = action      #action argument is a dictionary obj
        if  ("S" in actionobj.tags.keys() and actionobj.tags["S"][1] =="JavaScript") or ("SubType" in actionobj.tags.keys() and actionobj.tags["SubType"][1] =="JavaScript"):
            if "JS" in actionobj.tags.keys():
                print "\"JS\" in actionobj.tags.keys():"
                if actionobj.tags["JS"][0] =="refer":
                   print "action obj has a refer as /JS, refer is %s"%actionobj.tags["JS"][1]
                   return self.objects[actionobj.tags["JS"][1]].decode_stream()
                elif actionobj.tags["JS"][0] =="literal string":
                        print "get_javascript_from_actions() return:",actionobj.tags["JS"][1]
                        return actionobj.tags["JS"][1]
        return ""        

    def get_javascript_from_aditional_actions(self, aditional_action):
        if isinstance(aditional_action,str):    #aditional_action argument is the objnum
           aditional_action_obj = self.objects[aditional_action]
        else:
           aditional_action_obj = aditional_action      #aditional_action argument is a dictionary obj
        out = ""
        for key in aditional_action_obj.tags.keys():
            print "additional obj have a key:%s\n"%key
            if aditional_action_obj.tags[key][0] == "refer":
                out += self.get_javascript_from_actions(self.objects[aditional_action_obj.tags[key][1]])
            elif aditional_action_obj.tags[key][0] == "dictionary obj":
                out += self.get_javascript_from_actions(catalogobj.tags[key][1])
        return out  



 
    def get_header_from_annotation(self, annotation_obj):

        annotation_information = ""

        for key in annotation_obj.tags.keys():
            print "key %s in annotation_obj,type: %s"%(key,annotation_obj.tags[key][0])
            if key == "Subj":
               if annotation_obj.tags[key][0] =="refer":
                   annotation_information += 'annotation["subject"] = "%s";\n' % (self.objects[annotation_obj.tags[key][1]].decode_stream())
               elif annotation_obj.tags[key][0] =="literal string":
                   annotation_information += 'annotation["subject"] = "%s";\n' % (get_bracketed_content(annotation_obj.tags[key][1]))

        return annotation_information 
        


    @staticmethod
    def applyFilter(input):
        output = re.sub('^[\x00-\x19\x7f-\xff\n\s]*[\x00-\x19\x7f-\xff]','',input) #look for starting non-ascii characters
        output = re.sub('[\x00-\x19\x7f-\xff][\x00-\x19\x7f-\xff\r\s]*$','',output) #look for trailing non-ascii characters
        return output




def main(files):
    js_timeout=120
    pdf.DEBUG = True
    for file in files:
        print 'processing %s!!!' % file
        data = ''
        if os.path.exists(file):
            fin = open(file,'r')
            data = fin.read()
            fin.close()

        mypdf = pdf(data,file)
        if mypdf.is_valid():
            print 'parsing %s' % file
            mypdf.parse()

            mypdf.check_jbig2()

            decoded,decoded_headers = mypdf.getJavaScript()
            if len(decoded) > 0:
                decoded = decoded_headers + decoded
                fout = open(file+'.out','w')
                if fout:
                    fout_flag = True
                    print 'Wrote JavaScript (%d bytes -- %d headers / %d code) to file %s' % (len(decoded), len(decoded_headers), len(decoded) - len(decoded_headers), file+'.out') 
                    fout.write(decoded)
                    fout.close()
                    #execute JavaScript
                    t0 = time()
                    is_js_timeout=False
                    P = popen2.Popen3("./js -f pre.js -f %s -f post.js > %s"%(file+".out",file+".jsresult"), True)
                    while time() < t0 + js_timeout and P.poll() == -1:  
                         sleep(0.1)  
                    sts = P.poll()  
                    if sts == -1:  #timeout
                        is_js_timeout = True  
                        print "P.pid:",P.pid
                        os.kill(P.pid, signal.SIGTERM)        
                    p = P.fromchild
                    e = P.childerr
                    w = P.tochild
                    out = p.read()
                    err = e.read()
                    t1 = time()
                    print "t1",t1
                    if len(err.strip()) != 0: 
                        errorf = open(file+".jserror", 'w')
                        errorf.write(err)
                        errorf.close()
                    if is_js_timeout:
                        f = open(file+".jstimeout", 'w')
                        f.write('time_%s' % (t1-t0))
                        f.close()
                    else:
                        f = open(file+".jstime", 'w')
                        f.write('time_%s' % (t1-t0))
                        f.close()
                    p.close()
                    w.close()
                    e.close()
            else:
                print 'Didnt decode any JavaScript within PDF file'
        else:
            print('warn: ignoring non-pdf file ' + file)


if __name__ == '__main__':
    for i in sys.argv[1:]:
        main(glob.glob(i))
